import math
from .reward_function_base import BaseRewardFunction
from ..core.catalog import Catalog as c


class HeadingReward(BaseRewardFunction):
    """
    测量当前航向和目标航向之间的差异
    """
    def __init__(self, config):
        super().__init__(config)
        self.reward_item_names = [self.__class__.__name__ + item for item in ['', '_heading', '_alt', '_roll', '_speed']]
    
    def get_reward(self, task, env, agent_id):
        """
        将奖励建模为多个变量的高斯分布的乘积，然后取几何平均值来计算奖励
        Args:
            task: task instance
            env: environment instance

        Returns:
            (float): reward
            
        1.定义了四个变量的比例参数,用于控制每个变量对奖励的影响:
            heading_error_scale 控制了偏航角误差对奖励的影响。
            alt_error_scale 控制了高度误差对奖励的影响。
            roll_error_scale 控制了滚转角误差对奖励的影响。
            speed_error_scale 控制了速度误差对奖励的影响。
        2.计算每个变量的奖励分数 (heading_r, alt_r, roll_r, speed_r)
            通过将每个变量的误差值除以相应的比例参数，然后将结果取负指数来建模高斯分布。
        3.计算最终的奖励值 reward
            它是四个奖励分数的几何平均值的四次方根。这种方式可以确保奖励是正数，并且受到每个变量的影响。
        4.调用 _process 方法来处理奖励值，将四个奖励分数作为渲染项传递
        
        总之,get_reward 方法使用了四个不同的变量（偏航角误差、高度误差、滚转角误差和速度误差）来计算奖励值，
        并通过将它们的高斯分布的乘积取几何平均值来生成最终的奖励值。这个方法的目的是根据这些误差来鼓励代理在这些方面表现良好，奖励值越高，代表代理的性能越好。
        这种方式可以帮助代理在这些方面进行优化。
        """
        heading_error_scale = 5.0  # degrees
        heading_r = math.exp(-((env.agents[agent_id].get_property_value(c.delta_heading) / heading_error_scale) ** 2))

        alt_error_scale = 15.24  # m
        alt_r = math.exp(-((env.agents[agent_id].get_property_value(c.delta_altitude) / alt_error_scale) ** 2))

        roll_error_scale = 0.35  # radians ~= 20 degrees
        roll_r = math.exp(-((env.agents[agent_id].get_property_value(c.attitude_roll_rad) / roll_error_scale) ** 2))

        speed_error_scale = 24  # mps (~10%)
        speed_r = math.exp(-((env.agents[agent_id].get_property_value(c.delta_velocities_u) / speed_error_scale) ** 2))

        reward = (heading_r * alt_r * roll_r * speed_r) ** (1 / 4)
        return self._process(reward, agent_id, (heading_r, alt_r, roll_r, speed_r))